# -*- coding: utf-8 -*-
import random
import hashlib
import constants
import array
import struct


class StringCaseMutator:
    """
    本类提供三种变异方法：mutationLowerCase（小写变异）, mutationUpperCase（大写变异）, mutationRandomCase（随机大小写变异）
    """

    def __init__(self):
        self.name = "StringMutations"
        self._mutations = [
            self.mutationLowerCase,
            self.mutationUpperCase,
            self.mutationRandomCase,
        ]
        self._count = len(self._mutations)
        self._index = 0

    def next(self):
        """
        该函数用于序列变异，每次使_index+1，返回0代表序列变异完成
        :return:0表示到达最大值，1表示正常
        """
        self._index += 1
        if self._index >= self._count:
            return 0
        return 1

    def mutationLowerCase(self, data):
        return data.lower()

    def mutationUpperCase(self, data):
        return data.upper()

    def mutationRandomCase(self, data):
        # Allow us to skip ahead and always get same number
        rand = random.Random()
        rand.seed(hashlib.sha512(str(self._count)).digest())
        if len(data) > 20:
            for i in rand.sample(range(len(data)), 20):
                c = data[i]
                c = rand.choice([c.lower(), c.upper()])
                data = data[:i] + c + data[i + 1:]
            return data
        for i in range(len(data)):
            c = data[i]
            c = rand.choice([c.lower(), c.upper()])
            data = data[:i] + c + data[i + 1:]
        return data


class UnicodeStringsMutator:
    """
    利用constants.py中的UnicodeStringsMutator数据集对字符串进行变异
    """
    values = constants.UnicodeStringsMutator  # 得到unicode字符串（存放在constants.py文件中）

    def __init__(self):
        self.name = "UnicodeStringsMutator"
        if UnicodeStringsMutator.values is None:    # 如果文件中没有，就用_genValues自己生成
            print("Initialize %s" % self.name)
            self._genValues()
        self._maxCount = len(self.values)
        self._count = 0

    def _genValues(self):   # 生成unicode字符串
        if UnicodeStringsMutator.values is None:
            values = []

            sample = random.sample(range(2, 6024), 200)
            sample.append(0)
            sample.append(1024 * 65)

            uchars = range(0, 0xffff)

            for length in sample:
                value = u""
                for i in range(length):
                    value += unichr(random.choice(uchars))
                values.append(value)

            UnicodeStringsMutator.values = values

    def next(self):
        """
        Go to next mutation. When this is called the state machine is updated
        as needed.
        """
        self._count += 1
        if self._count >= self._maxCount:
            self._count -= 1
            return 0
        return 1

    # 利用count进行计数，来进行序列变异（从unicode中一个个选数据）
    def sequentialMutation(self):
        return self.values[self._count]

    # 利用choice进行随机选择变异
    def randomMutation(self, rand):
        return rand.choice(self.values)


class UnicodeBomMutator:
    """
    使用bom进行变异，bom就是下面那东西
    """
    values = constants.UnicodeBomMutator
    boms = ['\xFE\xFF', '\xFF\xEF', '\xEF\xBB\xBF'] # 这个应该就是bom了。。。

    def __init__(self):
        self.name = "UnicodeBomMutator"
        if UnicodeBomMutator.values is None:  # 这里就是None，所以加上了bom
            print("Initialize %s" % self.name)
            self._genValues()
        self._count = 0
        self._maxCount = len(self.values)

    def _genValues(self):
        if UnicodeBomMutator.values is None:
            valuesWithBOM = []
            values = []
            sample = random.sample(range(2, 2024, 2), 200)
            sample.append(0)
            sample.append(1024 * 2)
            for r in sample:
                values.append('A' * r)
            # 1. Prefix with both BOMs
            for v in values:
                for b in self.boms:
                    valuesWithBOM.append(b + v)
            # 2. Every other wchar
            for v in values:
                for b in self.boms:
                    newval = b
                    for i in range(0, len(v), 2):
                        newval += v[i:i + 2]
                        newval += b
                    valuesWithBOM.append(newval)
            # 3. Just BOM's
            for r in sample:
                newval = ""
                for i in range(r):
                    newval += random.choice(self.boms)
                valuesWithBOM.append(newval)
            UnicodeBomMutator.values = valuesWithBOM
            values = None

    def next(self):
        """
        调用该方法可以使count+1，直到达到maxcount（unicode字符串的最大值）
        """
        self._count += 1
        if self._count >= self._maxCount:
            self._count -= 1
            return 0
        return 1

    def getCount(self):
        return self._maxCount

    def sequentialMutation(self):
        return self.values[self._count]

    def randomMutation(self,  rand):
        return rand.choice(self.values)


class UnicodeBadUtf8Mutator:
    """
    Generate bad UTF-8 strings.
    """

    values = constants.UnicodeBadUtf8Mutator  # 这个也是None

    def __init__(self):
        self.name = "UnicodeBadUtf8Mutator"
        if UnicodeBadUtf8Mutator.values is None:    # 这个就是None
            print("Initialize %s" % self.name)
            self._genValues()
        self._count = 0
        self._maxCount = len(self.values)

    def _genValues(self):
        if UnicodeBadUtf8Mutator.values is None:
            encoding = [
                self._utf8OneByte,
                self._utf8TwoByte,
                self._utf8ThreeByte,
                self._utf8FourByte,
                self._utf8FiveByte,
                self._utf8SixByte,
                self._utf8SevenByte
            ]
            endValues = []
            sample = random.sample(range(2, 2024, 2), 100)
            sample.append(2024)
            for s in random.sample(range(2, 100), 50):
                sample.append(s)
            ascii = range(32, 126)
            for r in sample:
                value = ""
                for i in range(r):
                    value += random.choice(encoding)(random.choice(ascii))
                endValues.append(value)
            UnicodeBadUtf8Mutator.values = endValues

    '''
    FORMAT	C TYPE	PYTHON TYPE	STANDARD SIZE	NOTES
    x	pad byte	no value	 	 
    c	char	string of length 1	1	 
    b	signed char	integer	1	(3)
    B	unsigned char	integer	1	(3)
    ?	_Bool	bool	1	(1)
    h	short	integer	2	(3)
    H	unsigned short	integer	2	(3)
    i	int	integer	4	(3)
    I	unsigned int	integer	4	(3)
    l	long	integer	4	(3)
    L	unsigned long	integer	4	(3)
    q	long long	integer	8	(2), (3)
    Q	unsigned long long	integer	8	(2), (3)
    f	float	float	4	(4)
    d	double	float	8	(4)
    s	char[]	string	 	 
    p	char[]	string	 	 
    P	void *	integer	 	(5), (3)
    '''
    # struct中的pack函数按照给定的格式将数据封装成字符串，下面方法参数c，指的就是格式
    def _utf8OneByte(self, c):
        return struct.pack("!B", c)

    def binaryFormatter(self, num, bits=None, strip=False):
        if bits is None:
            bits = 64
            strip = True
        if type(num) == str:
            raise Exception("Strings not permitted")
        ret = ""
        for i in range(bits - 1, -1, -1):
            ret += str((num >> i) & 1)
        if strip:
            return ret.lstrip('0')
        return ret

    def _utf8TwoByte(self, c, mask='1100000010000000'):
        bfC = self.binaryFormatter(c)
        if len(bfC) > 11:
            raise Exception("Larger than two byte UTF-8")
        bfC = self.binaryFormatter(c, 11)
        bf = array.array('c', mask)
        bf[3:8] = array.array('c', bfC[0:5])
        bf[10:16] = array.array('c', bfC[5:])
        bfs = bf.tostring()
        return struct.pack("!BB", int(bfs[0:8], 2), int(bfs[8:16], 2))



    def _utf8ThreeByte(self, c, mask='111000001000000010000000'):
        bfC = self.binaryFormatter(c)
        if len(bfC) > 16:
            raise Exception("Larger than three byte UTF-8")
        bfC = self.binaryFormatter(c, 16)
        bf = array.array('c', mask)
        bf[4:8] = array.array('c', bfC[:4])
        bf[10:16] = array.array('c', bfC[4:10])
        bf[18:24] = array.array('c', bfC[10:])
        bfs = bf.tostring()
        return struct.pack("!BBB",
                           int(bfs[0:8], 2),
                           int(bfs[8:16], 2),
                           int(bfs[16:24], 2))



    def _utf8FourByte(self, c, mask='11110000100000001000000010000000'):
        bfC = self.binaryFormatter(c)
        if len(bfC) > 21:
            raise Exception("Larger than four byte UTF-8")
        bfC = self.binaryFormatter(c, 21)
        bf = array.array('c', mask)
        bf[5:8] = array.array('c', bfC[:3])
        bf[10:16] = array.array('c', bfC[3:9])
        bf[18:24] = array.array('c', bfC[9:15])
        bf[26:32] = array.array('c', bfC[15:])
        bfs = bf.tostring()
        return struct.pack("!BBBB",
                           int(bfs[0:8], 2),
                           int(bfs[8:16], 2),
                           int(bfs[16:24], 2),
                           int(bfs[24:32], 2))
    

    def _utf8FiveByte(self, c, mask='1111100010000000100000001000000010000000'):
        bfC = self.binaryFormatter(c)
        if len(bfC) > 26:
            raise Exception("Larger than five byte UTF-8")
        bfC = self.binaryFormatter(c, 26)
        bf = array.array('c', mask)
        bf[6:8] = array.array('c', bfC[:2])
        bf[10:16] = array.array('c', bfC[2:8])
        bf[18:24] = array.array('c', bfC[8:14])
        bf[26:32] = array.array('c', bfC[14:20])
        bf[34:40] = array.array('c', bfC[20:])
        bfs = bf.tostring()
        return struct.pack("!BBBBB",
                           int(bfs[0:8], 2),
                           int(bfs[8:16], 2),
                           int(bfs[16:24], 2),
                           int(bfs[24:32], 2),
                           int(bfs[32:40], 2))

    def _utf8SixByte(self, c, mask='111111001000000010000000100000001000000010000000'):
        bfC = self.binaryFormatter(c)
        if len(bfC) > 31:
            raise Exception("Larger than six byte UTF-8")
        bfC = self.binaryFormatter(c, 31)
        bf = array.array('c', mask)
        bf[7] = bfC[0]
        bf[10:16] = array.array('c', bfC[1:7])
        bf[18:24] = array.array('c', bfC[7:13])
        bf[26:32] = array.array('c', bfC[13:19])
        bf[34:40] = array.array('c', bfC[19:25])
        bf[42:48] = array.array('c', bfC[25:31])
        bfs = bf.tostring()
        return struct.pack("!BBBBBB",
                           int(bfs[0:8], 2),
                           int(bfs[8:16], 2),
                           int(bfs[16:24], 2),
                           int(bfs[24:32], 2),
                           int(bfs[32:40], 2),
                           int(bfs[40:48], 2))

    def _utf8SevenByte(self, c, mask='11111110100000001000000010000000100000001000000010000000'):
        bfC = self.binaryFormatter(c, 36)
        bf = array.array('c', mask)
        bf[10:16] = array.array('c', bfC[:6])
        bf[18:24] = array.array('c', bfC[6:12])
        bf[26:32] = array.array('c', bfC[12:18])
        bf[34:40] = array.array('c', bfC[18:24])
        bf[42:48] = array.array('c', bfC[24:30])
        bf[50:56] = array.array('c', bfC[30:])
        bfs = bf.tostring()
        return struct.pack("!BBBBBBB",
                           int(bfs[0:8], 2),
                           int(bfs[8:16], 2),
                           int(bfs[16:24], 2),
                           int(bfs[24:32], 2),
                           int(bfs[32:40], 2),
                           int(bfs[40:48], 2),
                           int(bfs[48:], 2))

    def next(self):
        """
        Go to next mutation. When this is called the state machine is updated
        as needed.
        """
        self._count += 1
        if self._count >= self._maxCount:
            self._count -= 1
            return 0
        return 1

    def getCount(self):
        return self._maxCount + 1

    def sequentialMutation(self):
        return self.values[self._count]

    def randomMutation(self, rand):
        return rand.choice(self.values)


class UnicodeUtf8ThreeCharMutator(UnicodeBadUtf8Mutator):
    """
    Generate long UTF-8 three byte strings
    """

    values = constants.UnicodeUtf8ThreeCharMutator

    def __init__(self):
        UnicodeBadUtf8Mutator.__init__(self)
        self.name = "UnicodeUtf8ThreeCharMutator"

        if UnicodeUtf8ThreeCharMutator.values is None:
            print("Initialize %s" % self.name)
            self._genValues()

    def _genValues(self):
        if UnicodeUtf8ThreeCharMutator.values is None:
            endValues = []
            sample = random.sample(range(2, 1024, 2), 300)
            sample.append(1024)
            # 1. Three char encoded values (can cause odd overflows)
            for r in sample:
                s = self._utf8ThreeByte(0xf0f0)
                endValues.append(s * r)
            UnicodeUtf8ThreeCharMutator.values = endValues


class StringMutator:
    """
    Apply StringFuzzer to each string node in DDT one Node at a time.
    """

    values = constants.StringMutator

    def __init__(self):
        self.name = "StringMutator"
        if StringMutator.values is None:
            print("Initialize {}".format(self.name))
            self._genValues()
        self._count = 0
        self._maxCount = len(self.values)

    def _genValues(self):
        values = []
        sample = random.sample(range(0, 6024), 200)  # 从0-6024序列中选择200个随机数
        sample.append(1024 * 65)
        uchars = range(0, 0xffff)
        for length in sample:
            value = u""
            for _ in range(length):
                value += unichr(random.choice(uchars))
            values.append(value)
        StringMutator.values = values

    def next(self):
        """
        Go to next mutation. When this is called the state machine is updated
        as needed.
        """
        self._count += 1
        if self._count >= self._maxCount:
            self._count -= 1
            return 0
        return 1

    def getCount(self):
        return self._maxCount

    def sequentialMutation(self):
        return self.values[self._count]

    def randomMutation(self,  rand):
        return rand.choice(self.values)

